from typing import Optional, Dict, Any
from openai import OpenAI
from config import logger
import json

client = OpenAI(
    api_key="sk-e9fc8cde99dc47f8970d514dd7d941f2", base_url="https://api.deepseek.com"
)


class LLMClient:
    @staticmethod
    def call_llm(prompt: str) -> Optional[str]:
        try:
            response = client.chat.completions.create(
                model="deepseek-chat",
                messages=[{"role": "user", "content": prompt}],
                timeout=300,
            )
            logger.error(response)
            return response.choices[0].message.content
        except Exception as e:
            logger.error(f"LLM API调用失败:{str(e)}")
            return None

    @staticmethod
    def call_llm_json(query: str) -> Optional[Dict[str, Any]]:
        content = LLMClient.call_llm(query)
        if not content:
            return None
        try:
            if "```json" in content:
                content = content.split("```json")[1].split("```")[0].strip()
                return json.loads(content)
            elif "```" in content:
                content = content.split("```")[1].split("```")[0].strip()
                return json.loads(content)
            return json.loads(content)
        except json.JSONDecodeError as e:
            logger.error(f"JSON解析失败:{str(e)}")
            return None

    @staticmethod
    def call_llm_sql(query: str) -> Optional[Dict[str, Any]]:
        content = LLMClient.call_llm(query)
        if not content:
            return None
        try:
            if "```sql" in content:
                content = content.split("```sql")[1].split("```")[0].strip()
                return content
            elif "```" in content:
                content = content.split("```")[1].split("```")[0].strip()
                return content
            return content
        except Exception as e:
            logger.error(f"SQL解析失败:{str(e)}")
            return None
